#!/usr/bin/env python
from __future__ import with_statement
"""
    parse_grape_results.py

"""

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   options        


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

import sys, os
import os.path
import StringIO

# add self to search path for testing
if __name__ == '__main__':
    exe_path = os.path.split(os.path.abspath(sys.argv[0]))[0]
    sys.path.append(os.path.abspath(os.path.join(exe_path,"..", "python_modules")))
    myname = os.path.split(sys.argv[0])[1]
    myname = os.path.splitext(myname)[0];
else:
    myname = __name__



#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   imports        


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
from itertools import izip
import sys, os


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   module code


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
#
#   grape output files come in sets of 5 lines
#       header, seq1, seq2, alignment_score, space
# 
HEADER  = 0
SEQ_1   = 1
SEQ_2   = 2
SCORE   = 3
GRAPE_LINES_BLOCK = 5

#_________________________________________________________________________________________            

#   get_sequence_ids_aligned_by_grape

#_________________________________________________________________________________________            
def get_sequence_ids_aligned_by_grape (grape_result_file_names):
    """
    Goes through each grape output file and finds out what sequences have been aligned
    """

    #
    #   grape output files come in sets of 5 lines
    #       header, seq1, seq2, alignment_score, space
    # 
    HEADER            = 0
    GRAPE_LINES_BLOCK = 5

    seq_ids = set()
    for axt_file in grape_result_file_names:
        for i, line in enumerate(open(axt_file)):
            #
            #   only interested in header
            # 
            if i % GRAPE_LINES_BLOCK != HEADER:
                continue

            # parse axt header ignore alignment index [0] and scores [7-]
            id1, beg1, end1, id2, beg2, end2, strand = line.split(" ", 8)[1:8]

            seq_ids.add(id1)
            seq_ids.add(id2)

    return seq_ids

    
        
#_________________________________________________________________________________________            

#   check_sequence_id

#_________________________________________________________________________________________            
def check_sequence_id (seq_id, seq_ids, stream_name, line_num, offending_line):
    """
    try sequence names with or without "chr" prefix
    """
    if seq_id in seq_ids:
        return seq_id
    
    if seq_id[3:] in seq_ids:
        return seq_id[3:]
    
    # sequence identifier not found: complain loudly
    raise Exception ("Unknown sequence identifier " +
                    "[%s] in file [%s:%d] [%s]" % ( seq_id, 
                                                    stream_name, 
                                                    line_num, offending_line))

#_________________________________________________________________________________________            

#   grape_same_scores_segments

#_________________________________________________________________________________________            
def grape_same_scores_segments(id1, id2, 
                                    saved_lines,
                                    beg, 
                                    end,
                                    strand, 
                                    beg1, end1, beg2, end2):
    """
    Goes through each aligned segment and breaks up into
    ungapped loci with equal posterior probabilities
    ignoring gap positions
    """

    seq1, seq2, scores = saved_lines[SEQ_1 : SCORE + 1]
    curr_pos1  = beg1                               
    curr_pos2  = beg2 if strand else end2 - beg
    #
    #  beginning of current segment
    # 
    curr_start1 = 0
    curr_start2 = 0
    prev_gap    = True
    prev_score = ''
    seq2_inc = 1 if strand else -1
    for score, s1, s2 in izip(scores[beg:end], seq1[beg:end], seq2[beg:end]):
        curr_gap = s1 == '-' or s2 == '-'
        #
        # changed state from one score to another or from a gapped to an ungapped state 
        # 
        if score != prev_score or curr_gap != prev_gap:
            #
            #   yield coordinates if we are in the middle of a segment
            # 
            if not prev_gap:
                if strand:
                    yield(id1, id2, strand, prev_score, curr_start1, curr_pos1, curr_start2, curr_pos2)
                else:
                    yield(id1, id2, strand, prev_score, curr_start1, curr_pos1, curr_pos2, curr_start2)    
            curr_start1 = curr_pos1
            curr_start2 = curr_pos2
        if s1 != '-':
            curr_pos1 += 1
        if s2 != '-':
            curr_pos2 += seq2_inc
        prev_gap = curr_gap
        prev_score = score
    if not prev_gap:
        if strand:
            yield(id1, id2, strand, prev_score, curr_start1, curr_pos1, curr_start2, curr_pos2)
        else:
            yield(id1, id2, strand, prev_score, curr_start1, curr_pos1, curr_pos2, curr_start2)    

#_________________________________________________________________________________________            

#   grape_aligned_positions

#_________________________________________________________________________________________            
def grape_aligned_positions(id1, id2, 
                            saved_lines,
                            beg, 
                            end,
                            strand, 
                            beg1, end1, beg2, end2):
    """
    Goes through each aligned segment and ignores gapped positions
    """

    seq1, seq2, scores = saved_lines[SEQ_1 : SCORE + 1]
    curr_pos1  = beg1                               
    curr_pos2  = beg2 if strand else end2 - beg - 1
    #
    #  beginning of current segment
    # 
    seq2_inc = 1 if strand else -1
    for score, s1, s2 in izip(scores[beg:end], seq1[beg:end], seq2[beg:end]):
        curr_gap = s1 == '-' or s2 == '-'
        if not curr_gap:
            yield(id1, id2, strand, score, curr_pos1, curr_pos2)
        if s1 != '-':
            curr_pos1 += 1
        if s2 != '-':
            curr_pos2 += seq2_inc

#_________________________________________________________________________________________            

#   grape_aligned_ends

#_________________________________________________________________________________________            
def grape_aligned_ends(id1, id2, 
                            saved_lines,
                            beg, 
                            end,
                            strand, 
                            beg1, end1, beg2, end2):
    """
    Goes through each aligned segment and ignores gapped positions
    """

    seq1, seq2, scores = saved_lines[SEQ_1 : SCORE + 1]
    curr_pos1  = beg1                               
    curr_pos2  = beg2 if strand else end2 - beg - 1
    #
    #  beginning of current segment
    # 
    seq2_inc = 1 if strand else -1
    


    cnt_gaps1 = seq1.count('-', beg, end)
    cnt_gaps2 = seq2.count('-', beg, end)
    end_pos1 = curr_pos1 + (end - beg - cnt_gaps1)
    end_pos2 = curr_pos2 + (end - beg - cnt_gaps2) * seq2_inc
    
    return (id1, id2, strand, scores[beg], curr_pos1, curr_pos2, scores[end-1], end_pos1, end_pos2)


#_________________________________________________________________________________________            

#   parse_grape_result_stream

#_________________________________________________________________________________________            
def parse_grape_result_stream (stream, stream_name, sequences, seq_ids):
    """
    Goes through stream and yields parameters for each aligned segment
    """

    saved_lines = [None, None, None, None, ""]

    for line_num, line in enumerate(stream):
        line_type = line_num % GRAPE_LINES_BLOCK

        #
        #   only need to do something after we have read the scores line
        # 
        saved_lines[line_type] = line.strip()
        if line_type != SCORE:
            continue


        #
        #   check that the lines are the same length
        #             
        if (len(saved_lines[SEQ_1]) != len(saved_lines[SEQ_2]) or 
            len(saved_lines[SEQ_1]) != len(saved_lines[SCORE])):
            raise Exception("Line lengths in file [%s:%d] do not match [%s]" % 
                                (stream_name, line_num - 2, saved_lines[HEADER]))


        # parse axt header ignore alignment index [0] and scores [7-]
        id1, beg1, end1, id2, beg2, end2, strand = saved_lines[0].split(" ", 8)[1:8]

        #
        # make sure sequence identifiers are correct
        id1 = check_sequence_id (id1, seq_ids, stream_name, line_num, saved_lines[HEADER])
        id2 = check_sequence_id (id2, seq_ids, stream_name, line_num, saved_lines[HEADER])


        strand = strand == "+"
        
        # convert from string to integer
        beg1, end1, beg2, end2= map(int, (beg1, end1, beg2, end2))

        # 
        # count from 0 not from 1
        # 
        beg1 -= 1  
        beg2 -= 1

        #
        #   grape will align segments against "virtual" sequence if the posterior
        #       probability scores are high, 
        #       i.e. extending the alignment by adding initial and terminal 'N's 
        #
        #   Ignore this by 1) adjusting for -ve offsets
        #   Ignore this by 2) adjusting for sequence length
        # 

        beg_virtual_len = 0
        end_virtual_len = 0
        seq_len1 = len(sequences[id1])
        seq_len2 = len(sequences[id2])

        if beg1 < 0:
            beg_virtual_len = - beg1
        if end1 > seq_len1:
            end_virtual_len = end1 - seq_len1
        if beg2 < 0:
            if strand:
                assert(beg_virtual_len == 0)
                beg_virtual_len = - beg2
            else:
                assert(end_virtual_len == 0)
                end_virtual_len = - beg2
        if end2 > seq_len2:
            if strand:
                assert(end_virtual_len == 0)
                end_virtual_len = end2 - seq_len2
            else:
                assert(beg_virtual_len == 0)
                beg_virtual_len = end2 - seq_len2
                                             
        beg1 += beg_virtual_len
        beg2 += beg_virtual_len
        align_end = len(saved_lines[SEQ_1])
        align_end -= end_virtual_len

        yield(      id1, id2, 
                    saved_lines,
                    beg_virtual_len, 
                    align_end,
                    strand, 
                    beg1, end1, beg2, end2)
        saved_lines = [None, None, None, None, ""]
    

#_________________________________________________________________________________________            

#   parse_grape_files_same_scores_segments

#_________________________________________________________________________________________            
def parse_grape_files_get_same_scores_segments (grape_result_file_names, sequences, seq_ids):
    """
    Goes through each grape output file and calls callback with aligned segments
    """

    #
    #   order sequences
    # 
    # 
    seq_ids = set(seq_ids)

    for grape_result_file_name in grape_result_file_names:
        for params in parse_grape_result_stream (   open(grape_result_file_name), 
                                                    grape_result_file_name,
                                                    sequences, seq_ids):
            for (id1, id2, strand, score,
                 b1, e1, b2, e2) in grape_same_scores_segments(*params):
                yield (id1, id2, strand, score, b1, e1, b2, e2)
        

#_________________________________________________________________________________________            

#   parse_grape_files_same_scores_segments

#_________________________________________________________________________________________            
def parse_grape_files_get_aligned_positions (grape_result_file_names, sequences, seq_ids):
    """
    Goes through each grape output file and calls callback with aligned segments
    """

    #
    #   order sequences
    # 
    # 
    seq_ids = set(seq_ids)

    for grape_result_file_name in grape_result_file_names:
        for params in parse_grape_result_stream (   open(grape_result_file_name), 
                                                    grape_result_file_name,
                                                    sequences, seq_ids):
            for (id1, id2, strand, score,
                 p1, p2) in grape_aligned_positions(*params):
                yield (id1, id2, strand, score, p1, p2)

#_________________________________________________________________________________________            

#   parse_grape_files_same_scores_segments

#_________________________________________________________________________________________            
def parse_grape_files_alignment_ends (grape_result_file_names, sequences, seq_ids):
    """
    Goes through each grape output file and calls callback with aligned segments
    """

    #
    #   order sequences
    # 
    # 
    seq_ids = set(seq_ids)

    for grape_result_file_name in grape_result_file_names:
        for params in parse_grape_result_stream (   open(grape_result_file_name), 
                                                    grape_result_file_name,
                                                    sequences, seq_ids):
            yield grape_aligned_ends(*params)


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Unit testing


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

import unittest
if __name__ == '__main__':
    from collections import defaultdict
    from custom_log import custom_log
    from fasta_sequences_iterator import fasta_sequences_iterator

    
class TestParseGrapeResults(unittest.TestCase):

    #       self.assertEqual(self.seq, range(10))
    #       self.assert_(element in self.seq)
    #       self.assertRaises(ValueError, random.sample, self.seq, 20)



    def setUp(self):
        """
            use data in unit_testing_data for unit testing
        """
        self.logger = custom_log(myname, "parse_grape_results.log")

    def test_get_sequence_ids_aligned_by_grape(self):
        """
            check get_sequence_ids_aligned_by_grape
        """
        ids = get_sequence_ids_aligned_by_grape ([os.path.join(exe_path, "unit_testing_data", "test.grape_results")])
        self.assert_(ids == set(['test1', 'test2','test3', 'test4']))
        
        
        
    def test_parse_grape_result_stream (self):
        #
        #   load test sequences
        # 
        seq_file_name = os.path.join(exe_path, "unit_testing_data", "test.sequence.fa")
        from fasta_sequences_iterator import fasta_sequences_iterator
        sequences = dict()
        for seq_id, seq in fasta_sequences_iterator(open(seq_file_name)):
            sequences[seq_id] = seq
        
                    
        #
        #   grape results
        # 
        results_file_name = os.path.join(exe_path, "unit_testing_data", "test.grape_results")

        aligned_segments1 = defaultdict(list)
        aligned_segments2 = defaultdict(list)
        aligned_positions1 = defaultdict(list)
        aligned_positions2 = defaultdict(list)
        for params in parse_grape_result_stream (open(results_file_name), 
                                                    results_file_name, sequences, 
                                                    sequences.keys()):
            (id1, id2, alignment,
             beg, end, strand, 
             beg1, end1, beg2, end2) = params
            # Make sure unaligned "virtual sequence" off the ends of the 
            #   sequences is handled correctly
            seq1, seq2  = alignment[1:3]
            self.assert_(seq1[end:] == "unalignable"  )
            self.assert_(seq2[end:] == "nnnnnnnnnnn"  )
            self.assert_(seq2[:beg] == "naligned")
            self.assert_(seq1[:beg] == "NNNNNNNN")

            
            #
            #   save parsed aligned_segments and make sure they are correct
            # 
            for (id1, id2, strand, score,
                 b1, e1, b2, e2) in grape_same_scores_segments(*params):
                if strand:
                    aligned_segments1[score].append(sequences[id1][b1:e1])
                    aligned_segments2[score].append(sequences[id2][b2:e2])
                else:
                    aligned_segments1[score].append(sequences[id1][b1:e1])
                    aligned_segments2[score].append(sequences[id2][b2:e2][::-1])

            #
            #   save parsed aligned_segments and make sure they are correct
            # 
            for (id1, id2, strand, score,
                 p1, p2) in grape_aligned_positions(*params):
                aligned_positions1[score].append(sequences[id1][p1])
                aligned_positions2[score].append(sequences[id2][p2])

        #
        #   make sure segments are correctly parsed 
        #
        self.assert_(aligned_segments1 == 
                                defaultdict(list,
                                            {'A': ['DEFGHIJ', 'DEFGHIJ'], 
                                             'B': ['L', 'L'], 
                                             'C': ['NOPQRST', 'NOPQRST'], 
                                             'D': ['UVWX', 'UVWX'],
                                             'Z': ['ABC', 'K', 'YZ', 'ABC', 'K', 'YZ']
                                             }))
        #
        #   join up segments and individual positions. Should be the same 
        #
        for score in aligned_positions1:
            aligned_positions1[score] = "".join(aligned_positions1[score])
            aligned_positions2[score] = "".join(aligned_positions2[score])
            aligned_segments1[score] = "".join(aligned_segments1[score])
            aligned_segments2[score] = "".join(aligned_segments2[score])
        self.assert_(aligned_segments1, aligned_positions2)
        self.assert_(aligned_segments1, aligned_positions2)
        self.assert_(aligned_segments1, aligned_segments2)
            


                
#
#   debug code not run if called as a module
#     
if __name__ == '__main__':
    unittest.main()